import pandas as pd
import numpy as np


class DataLoader:
    """
    数据管线，支持行情、因子数据注入
    """

    def __init__(self, market_data: pd.DataFrame = None):
        """
        初始化DataLoader

        Parameters:
        market_data (pd.DataFrame): 行情数据，columns=['trading', 'code', 'close', 'open', 'high', 'low', 'volume']
            列trading: datetime.date格式
        factors_data: 因子数据，columns=[trading, code, *signals]
        """
        self.market_data = None
        self.factor_data = {}  # 存储因子数据的字典

        # 索引和缓存
        self.indexed_market_data = None
        self.grouped_market_by_date = None
        self.grouped_market_by_code = None
        self.unique_dates = None
        self.date_to_index = {}
        self._cache = {}  # 添加缓存机制

        if market_data is not None:
            self._preprocessing_marketData(market_data)

    def _preprocessing_marketData(self, data: pd.DataFrame):
        """
        预处理行情数据，生成数据索引
        """
        # 检查所有必要的列是否都存在于DataFrame中
        required_columns = ['trading', 'code', 'close', 'open', 'high', 'low', 'volume']
        missing_columns = [col for col in required_columns if col not in data.columns]
        assert not missing_columns, f"缺少必要的列: {missing_columns}"

        # 行情数据处理
        self.market_data = data.sort_values(by='trading', ascending=True, ignore_index=True)

        self._create_market_indexes()

        # 缓存唯一日期
        self.unique_dates = self.market_data['trading'].unique()
        self.date_to_index = {date: idx for idx, date in enumerate(self.unique_dates)}

    def _preprocessing_factors(self, data: pd.DataFrame, signal_name: list = None):
        """
        预处理因子数据
        * 支持动态新注入因子数据，通过signal_name参数区分
        """
        if signal_name is None:
            signal_name = []
        assert len(signal_name) > 0, "请传入因子名称"

        required_columns = ['trading', 'code'] + signal_name
        missing_columns = [col for col in required_columns if col not in data.columns]
        assert not missing_columns, f"缺少必要的列: {missing_columns}"

        indexed_factor_data = data.set_index(['trading', 'code'])

        for f in signal_name:
            _f = indexed_factor_data[f]
            self.factor_data[f] = _f

    def _create_market_indexes(self):
        if self.market_data is not None:
            # 创建多级索引
            self.indexed_market_data = self.market_data.set_index(['trading', 'code'])
            self.indexed_market_data.sort_index(inplace=True)

            # 按日期分组
            self.grouped_market_by_date = self.market_data.groupby('trading')

            # 按股票代码分组
            self.grouped_market_by_code = self.market_data.groupby('code')

    def get_close(self, date, code):
        # 根据日期和股票代码获取收盘价
        cache_key = ('close', date, code)
        if cache_key in self._cache:
            return self._cache[cache_key]
        try:
            result = self.indexed_market_data.loc[date, code]['close']
            if isinstance(result, pd.Series):
                result = result.values[0]
            if np.isnan(result):
                result = None
            self._cache[cache_key] = result
            return result
        except KeyError:
            return None

    def get_open(self, date, code):
        # 根据日期和股票代码获取开盘价
        cache_key = ('open', date, code)
        if cache_key in self._cache:
            return self._cache[cache_key]
        try:
            result = self.indexed_market_data.loc[date, code]['open']
            if isinstance(result, pd.Series):
                result = result.values[0]
            if np.isnan(result):
                result = None
            self._cache[cache_key] = result
            return result
        except KeyError:
            return None

    def get_volume(self, date, code):
        # 根据日期和股票代码获取成交量，使用缓存
        cache_key = ('volume', date, code)
        if cache_key in self._cache:
            return self._cache[cache_key]
        try:
            result = self.indexed_market_data.loc[date, code]['volume']
            if isinstance(result, pd.Series):
                result = result.values[0]
            if np.isnan(result):
                result = None
            self._cache[cache_key] = result
            return result
        except KeyError:
            return None

    def get_date(self, indexnum):
        # 根据索引获取日期
        if indexnum >= len(self.unique_dates):
            return None
        return self.unique_dates[indexnum]

    def clear_cache(self):
        # 清空缓存
        self._cache.clear()
